/**
 *  @file SSLContext.cpp
 */
#include "SSLContext.h"
#include "SSLException.h"
#include "SSLSocket.h"
#include "SSLServerSocket.h"
#include "../../security/__internal/InitOpenSSL.h"
#include "../../security/cert/X509Certificate.h"
#include "../../common/CppMemPool.h"
#include "../../common/ObjPtrT.h"


#if defined(linux)
#   include <stdio.h>
#endif

#include "openssl/ssl.h" 

using namespace cppflib;
using namespace cppflib::security::cert;



namespace cppflib
{

namespace net
{

namespace ssl
{


//! Maximum allowable length of a cert chain
PRIVATE const int st_defaultSSLVerifyDepth = 4;   
//! Data slot index for SSL object when calling callback to verify peer cert. See __ReadPeerCertCallback
PRIVATE volatile int st_slotForVerifyPeerCert = -1;


PRIVATE SSL_CTX *__GetSSLCtx(ptrany_t sslCtx)
{
   return static_cast<SSL_CTX*>(sslCtx);
}

/**
 *  Callback to provide password for private key
 */
PRIVATE int __PrivateKeyPwdCallback(char *buf, int size, int rwflag, void *userdata)
{
   int bufLen = 0;
   ISSLPrivateKeyPwdCallback *pCb = (userdata != NULL) ? reinterpret_cast<ISSLPrivateKeyPwdCallback*>(userdata) :
                                                         NULL;

   if (pCb != NULL) {
      CFString pwd;
      pCb->GetPwd(pwd);

      if (!pwd.IsEmpty()) {
         pwd.ToUtf8String(buf, size, &bufLen);
      }
   }

   return bufLen;
}

/**
 *  Get a slot index in order to embed a piece of userdata to SSL object
 */
PRIVATE int __GetSlotForVerifyPeerCertCallback()
{
   if (st_slotForVerifyPeerCert < 0) {
      st_slotForVerifyPeerCert = SSL_get_ex_new_index(0, NULL, NULL, NULL, NULL);
   }
   return st_slotForVerifyPeerCert;
}

/**
 *  Callback to examine cert chain sent from peer during negotiation
 *  If there are 3 certs in the cert chain, this function will be called 3 times.
 *
 *  @note this won't be called if peer doesn't provide a cert during negotiation
 *  
 *  @param [in] ok -- 0 means this cert is *NOT* OK
 *  @param [in] store -- store which contains the X509 cert
 */
PRIVATE int __ReadPeerCertCallback(int ok, X509_STORE_CTX *store) 
{ 

   if (st_slotForVerifyPeerCert >= 0) {
      SSL *sslObj = static_cast<SSL*>(X509_STORE_CTX_get_ex_data(store, SSL_get_ex_data_X509_STORE_CTX_idx()));
      ISSLPeerCertCallback *pCb = static_cast<ISSLPeerCertCallback*>(SSL_get_ex_data(sslObj, st_slotForVerifyPeerCert));
      X509 *cert = X509_STORE_CTX_get_current_cert(store);

      if (pCb != NULL && cert != NULL) {
         try {
            X509Certificate x509Cert(cert, false);
            pCb->ReadPeerCert((ok != 0), x509Cert);
         }
         catch(exceptions::BaseException &e) {
            e.GetErrCode();  // to eliminate compiler warning
            ok = 0; // don't want to continue
         }
      }
   }

   return ok;
}

/**
 *  A callback that will always trust the peer by returning '1'
 */
PRIVATE int __TrustPeerCertCallback(int ok, X509_STORE_CTX *store) 
{ 

   do {
      // since cert is not checked, the first time that this function is called will be passed with
      // a '0', which is redundant
      if (ok == 0) {
         ok = 1;
         break;
      }

      if (st_slotForVerifyPeerCert < 0)
         break;

      SSL *sslObj = static_cast<SSL*>(X509_STORE_CTX_get_ex_data(store, SSL_get_ex_data_X509_STORE_CTX_idx()));
      ISSLPeerCertCallback *pCb = static_cast<ISSLPeerCertCallback*>(SSL_get_ex_data(sslObj, st_slotForVerifyPeerCert));
      X509 *cert = X509_STORE_CTX_get_current_cert(store);

      if (!pCb || !cert)
         break;

      try {
         X509Certificate x509Cert(cert, false);
         pCb->ReadPeerCert((ok != 0), x509Cert);
      }
      catch(exceptions::BaseException &e) {
         e.GetErrCode(); // to eliminate compiler warning
         ok = 0; // don't want to continue
      }

   } while (false);

   return ok;  // always continue
}

/**
 *  Called when server negotiates with a client and client embeds a Server Name Indication (SNI) request.
 *  See SSLContext::SetServerNameCallback.
 *
 *  @param [in] ssObj -- the SSL object negotiating with client
 *  @param [in] ad -- not used
 *  @param [in] arg -- user provided data
 */
PRIVATE int __FindMatchingHostNameCallback(SSL *sslObj, int *ad, void *arg)
{
   do {
      const char * pServerName = SSL_get_servername(sslObj, TLSEXT_NAMETYPE_host_name);
      if (!pServerName)  // client doesn't embed a SNI
         break;

      if (!arg)  // user don't want to check
         break;

      CFString hostName(pServerName);
      ISSLFindMatchingHostNameCallback *pCb = static_cast<ISSLFindMatchingHostNameCallback*>(arg);
      SSLContext *pCtxt = pCb->GetMatchingSSLContext(hostName);
      if (!pCtxt)  // no matching hostname
         return SSL_TLSEXT_ERR_ALERT_WARNING;
      
      // switch context if they are different
      SSL_CTX * pOrgCtxt = SSL_get_SSL_CTX(sslObj);
      if (pOrgCtxt != __GetSSLCtx(pCtxt->GetInternalContext()))
         SSL_set_SSL_CTX(sslObj, __GetSSLCtx(pCtxt->GetInternalContext()));
      
   } while (false);
   
   return SSL_TLSEXT_ERR_OK;
}

//////////////////////////////// SSLContext begins ///////////////////////////////////////////

//! List of cipher to use. Those with a '!' at the front indicates 'not to use'
#define CIPHER_LIST "ALL:!ADH:!LOW:!EXP:!MD5:@STRENGTH"

/**
 *  Create a client side SSL context which will trust the server without checking the cert
 *  This is for testing only and should not be used in a real application
 */
SSLContext::SSLContext()
{
   cppflib::security::__internal::InitOpenSSL();

   SSL_CTX *ctx = SSL_CTX_new(SSLv23_method()); // include SSLv2, SSLv3 and TLSv1
   
   this->isServerCtx = false;
   this->sslCtx = reinterpret_cast<ptrany_t>(ctx);
   this->pPrivateKeyPwdCb = NULL;
   this->pFindMatchingHostNameCb = NULL;

   SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, &__TrustPeerCertCallback);  // set verify mode and callback
   SSL_CTX_set_verify_depth(ctx, st_defaultSSLVerifyDepth);   // set max length of cert chain

   // include workarounds for buggy peers and exclude SSLv2, which is insecure
   SSL_CTX_set_options(ctx, SSL_OP_ALL|SSL_OP_NO_SSLv2);  
   if (SSL_CTX_set_cipher_list(ctx, CIPHER_LIST) != 1) 
      THROW_CTOR_EXCEP(SSLException, _S("Setting cipher list failed (no valid ciphers)"));
}

/**
 *  Create a SSL context
 *
 *  @param [in] isServer -- a context for server or not
 *  @param [in] pathToPEMFileWithRootCerts -- PEM file which contains trusted root certs
 *  @param [in] needPeerAuth -- (for server context only) whether server requires peer (i.e. client) authentication
 *                              (client context always requires peer (server) to provide a certificate)
 */
SSLContext::SSLContext(bool isServer, const CFString &pathToPEMFileWithRootCerts, bool needPeerAuth)
{
   cppflib::security::__internal::InitOpenSSL();

   SSL_CTX *ctx = SSL_CTX_new(SSLv23_method()); // include SSLv2, SSLv3 and TLSv1
   ObjPtrT<char> pTempPEMPath(pathToPEMFileWithRootCerts.ToUtf8String());
   
   this->isServerCtx = isServer;
   this->sslCtx = reinterpret_cast<ptrany_t>(ctx);
   this->pPrivateKeyPwdCb = NULL;
   this->pFindMatchingHostNameCb = NULL;

   if (SSL_CTX_load_verify_locations(ctx, pTempPEMPath, NULL) != 1)  // load root certs
      THROW_CTOR_EXCEP(SSLException, _S("Loading root certs file failed"));
      //throw SSLException(_S("Loading root certs file failed")); 
   if (SSL_CTX_set_default_verify_paths(ctx) != 1)  // load certs from default location, too
      THROW_CTOR_EXCEP(SSLException, _S("Loading default root cert file failed"));

   int mode = SSL_VERIFY_PEER;
   if (isServer && needPeerAuth)
      mode |= SSL_VERIFY_FAIL_IF_NO_PEER_CERT;   // for server only

   SSL_CTX_set_verify(ctx, mode, &__ReadPeerCertCallback);  // set verify mode and callback
   SSL_CTX_set_verify_depth(ctx, st_defaultSSLVerifyDepth);   // set max length of cert chain

   //SSL_CTX_set_mode(ctx, SSL_MODE_AUTO_RETRY);

   // include workarounds for buggy peers and exclude SSLv2, which is insecure
   long ctxOptions = SSL_OP_ALL|SSL_OP_NO_SSLv2;
   if (isServer)
      ctxOptions |= SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
   SSL_CTX_set_options(ctx, ctxOptions);
   
   if (SSL_CTX_set_cipher_list(ctx, CIPHER_LIST) != 1) 
      THROW_CTOR_EXCEP(SSLException, _S("Setting cipher list failed (no valid ciphers)"));

   // set callback for SNI
   if (isServer) {
      SSL_CTX_set_tlsext_servername_callback(ctx, &__FindMatchingHostNameCallback);
      SSL_CTX_set_tlsext_servername_arg(ctx, NULL);
   }
}

SSLContext::~SSLContext(void)
{
   Close();
}

void SSLContext::Close()
{
   ReleaseBeforeThrowing();
}

void SSLContext::ReleaseBeforeThrowing()
{
   if (sslCtx) {
      SSL_CTX_free(__GetSSLCtx(sslCtx));
      sslCtx = NULL;
   }

   if (pPrivateKeyPwdCb) {
      cppmempool::Release(pPrivateKeyPwdCb);
      pPrivateKeyPwdCb = NULL;
   }

   if (pFindMatchingHostNameCb) {
      cppmempool::Release(pFindMatchingHostNameCb);
      pFindMatchingHostNameCb = NULL;
   }
   
   cppflib::security::__internal::UninitOpenSSL();
}


/**
 *  To set chain of certs which will be sent to peer. Client may or may not call this.
 *
 *  @param [in] pathToPEMFileWithCertChain -- path to file containing cert chain as well as the private key
 *  @param [in] pPrivateKeyPwdCb -- callback to provide password for the private key
 */
bool SSLContext::SetCertChain(const CFString &pathToPEMFileWithCertChain, ISSLPrivateKeyPwdCallback *pPrivateKeyPwdCb)
{
   if (!sslCtx) 
      return false;

   cppmempool::Retain(pPrivateKeyPwdCb);
   cppmempool::Release(this->pPrivateKeyPwdCb);
   this->pPrivateKeyPwdCb = pPrivateKeyPwdCb;

   // load chain of certificates
   ObjPtrT<char> pTempPEMCertFile(pathToPEMFileWithCertChain.ToUtf8String());
   if (SSL_CTX_use_certificate_chain_file(__GetSSLCtx(sslCtx), pTempPEMCertFile) != 1)
      throw SSLException(_S("Loading certificate chain failed"));

   // load private key
   SSL_CTX_set_default_passwd_cb(__GetSSLCtx(sslCtx), &__PrivateKeyPwdCallback);
   SSL_CTX_set_default_passwd_cb_userdata(__GetSSLCtx(sslCtx), reinterpret_cast<ptrany_t>(this->pPrivateKeyPwdCb));
   if (SSL_CTX_use_PrivateKey_file(__GetSSLCtx(sslCtx), pTempPEMCertFile, SSL_FILETYPE_PEM) != 1) 
      throw SSLException(_S("Loading private key failed"));

   return true;
}

/**
 *  Set callback to find SSLContext with matching hostname (for server socket only)
 */
bool SSLContext::SetServerNameCallback(ISSLFindMatchingHostNameCallback *pFindMatchingHostNameCb)
{
   if (!sslCtx || !isServerCtx) 
      return false;

   cppmempool::Retain(pFindMatchingHostNameCb);
   cppmempool::Release(this->pFindMatchingHostNameCb);
   this->pFindMatchingHostNameCb = pFindMatchingHostNameCb;
   SSL_CTX_set_tlsext_servername_arg(__GetSSLCtx(sslCtx), static_cast<ptrany_t>(this->pFindMatchingHostNameCb));

   return true;
}

/**
 *  Whether this SSL context requires peer to provide a certificate for authentication
 */
bool SSLContext::NeedPeerAuthentication()
{
   int verifyMode = SSL_CTX_get_verify_mode(__GetSSLCtx(sslCtx));

   if (this->isServerCtx) {
      // server context
      if (verifyMode & SSL_VERIFY_FAIL_IF_NO_PEER_CERT)
         return true;
   }
   else {
      // client context
      if (verifyMode & SSL_VERIFY_PEER)
         return true;
   }

   return false;
}
    
/**
 *  Create a client socket from this context (SSLServerSocket will also call this)
 *
 *  @param [in] pSSLSock -- Pass NULL if it is a client context.
 *                          For server context, it must be a valid instance.
 *
 *  @param [in] pPeerCertCb -- callback to peek at the certificates received. 
 *                             Can be NULL if not interested
 */
SSLSocket *SSLContext::CreateClientSocket(SSLSocket *pSSLSock, ISSLPeerCertCallback *pPeerCertCb)
{
   if (this->isServerCtx && pSSLSock == NULL)
      throw SSLException(_S("Server SSL context must provide a valid SSLSocket"));

   SSL *sslObj = SSL_new(__GetSSLCtx(sslCtx));
   if (!sslObj)
      throw SSLException(_S("Create SSL object failed"));

   // add callback to SSL object
   if (pPeerCertCb) 
      SSL_set_ex_data(sslObj, __GetSlotForVerifyPeerCertCallback(), static_cast<ptrany_t>(pPeerCertCb));

   SSLSocket *retSock = pSSLSock;
   if (retSock == NULL) {
       Socket::sockattr_t flag = NeedPeerAuthentication() ? 
                                    static_cast<Socket::sockattr_t>(Socket::SSL_CHK_PEERCERT | Socket::SSL_CHK_VERIFYRESULT) :
                                    static_cast<Socket::sockattr_t>(0);
       retSock = new SSLSocket(flag);
   }
   retSock->Initialize(sslObj, pPeerCertCb, this->isServerCtx);
   return retSock;
}

/**
 *  Create a server socket from this context
 *
 *  @param [in] pPeerCertCb -- callback to peek at the certificates received. 
 *                             Can be NULL if not interested
 */
SSLServerSocket *SSLContext::CreateServerSocket(ISSLPeerCertCallback *pPeerCertCb)
{
   if (!this->isServerCtx)
      throw SSLException(_S("Client SSL context can't create server socket"));

   SSLServerSocket *serverSock = new SSLServerSocket(this, pPeerCertCb);
   return serverSock;
}

/**
 *  Get the internal SSL context pointer
 */
ptrany_t SSLContext::GetInternalContext()
{
   return sslCtx;
}

/**
 *  Get OpenSSL library version
 */
CFString SSLContext::GetOpenSSLVersion()
{
   //
   //  Calling SSLeay() is better than using macro OPENSSL_VERSION_NUMBER directly
   //  since SSLeay() is embedded inside the library rather in the header file
   //
   u32_t ver = SSLeay();  
   u32_t major = (ver >> 28) & 0x0f;
   u32_t minor = (ver >> 20) & 0xff;
   u32_t fix = (ver >> 12) & 0xff;
   u32_t patch = (ver >> 4) & 0xff;
  
   char patchChar = (patch > 0) ? static_cast<char>(patch - 1 + 'a') : 0;
   CFString s;
   if (patchChar > 0)
      s.Format(_S("%d.%d.%d%c"), major, minor, fix, patchChar);
   else 
      s.Format(_S("%d.%d.%d"), major, minor, fix);
      
   return s;
}

/**
 *  Get OpenSSL library version in the following format
 *
 *  MNNFFPPS: major minor fix patch status
 */
u32_t SSLContext::GetOpenSSLVersionNumber()
{
   return SSLeay();
}

} // end of namespace ssl

} // end of namespace net

} // end of namespace cppflib
